"""
双向循环神经网络：mnist分类
"""

import os
import torch
import torch.nn as nn
import torchvision.datasets
import torch.utils.data
import torchvision.transforms as transforms
from tqdm import tqdm
import numpy as np
import matplotlib.pyplot as plt


def update_lr(optimizer, lr):
    """更新学习率"""
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr


class Config(object):
    """单例模式，配置参数"""
    _instance = None
    sequence_length = 28
    input_size = 28
    hidden_size = 128
    num_layers = 2
    num_classes = 10
    batch_size = 128
    num_epochs = 2
    nw = os.cpu_count()
    learning_rate = 0.001
    model_path = 'model/checkpoint/bi-rnn-1.pth'
    device_cuda = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    device_cpu = torch.device('cpu')

    def __new__(cls, *args, **kw):
        if cls._instance is None:
            cls._instance = object.__new__(cls, *args, **kw)
        return cls._instance


class Mnist:
    _instance = None
    config = Config()
    batch_size = config.batch_size
    nw = config.nw

    transform = transforms.Compose([transforms.ToTensor()])

    trainset = torchvision.datasets.MNIST(root='./dataset/mnist', train=True, download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(dataset=trainset, batch_size=batch_size, shuffle=True, num_workers=nw)

    testset = torchvision.datasets.MNIST(root='./dataset/mnist', train=False, download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(dataset=testset, batch_size=batch_size, shuffle=False, num_workers=nw)

    def __new__(cls, *args, **kw):
        if cls._instance is None:
            cls._instance = object.__new__(cls, *args, **kw)
        return cls._instance


class BiRNN(nn.Module):
    """双向循环神经网络"""

    def __init__(self, input_size, hidden_size, num_layers, num_classes, device=Config.device_cuda):
        super().__init__()
        self.device = device
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, bidirectional=True)
        self.fc = nn.Linear(hidden_size*2, num_classes)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers*2, x.size(0), self.hidden_size).to(self.device)
        c0 = torch.zeros(self.num_layers*2, x.size(0), self.hidden_size).to(self.device)
        out, _ = self.lstm(x, (h0, c0))
        out = self.fc(out[:, -1, :])
        return out


def train():
    """训练模型"""
    config = Config()
    device = config.device_cuda
    print(f'training on device: {device}')
    model_path = config.model_path
    model = BiRNN(config.input_size, config.hidden_size, config.num_layers, config.num_classes).to(device)
    if os.path.isfile(model_path):
        print('load model from: ', model_path)
        model.load_state_dict(torch.load(model_path))
        model.train()

    optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
    criterion = nn.CrossEntropyLoss()
    trainloader = Mnist.trainloader
    curr_lr = config.learning_rate
    for epoch in range(config.num_epochs):
        print(f'Epoch:{epoch + 1}/{config.num_epochs}')
        pbar = tqdm(trainloader)
        running_loss = []
        for images, labels in pbar:
            images = images.reshape(-1, config.sequence_length, config.input_size)
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            running_loss.append(loss.item())
            pbar.set_description_str(f'loss:{round(np.nanmean(running_loss), 3)}')

        if (epoch + 1) % 20 == 0:
            curr_lr /= 3
            update_lr(optimizer, curr_lr)

    # save model
    print(f'save model to: {model_path}')
    if not os.path.isdir(os.path.dirname(model_path)):
        os.makedirs(os.path.dirname(model_path))
    torch.save(model.state_dict(), model_path)


def test():
    """测试模型"""
    config = Config()
    model_path = config.model_path
    device = config.device_cuda
    model = BiRNN(config.input_size, config.hidden_size, config.num_layers, config.num_classes).to(device)
    if os.path.isfile(model_path):
        print('load model from: ', model_path)
        model.load_state_dict(torch.load(model_path))

    model.eval()
    testloader = Mnist.testloader
    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in testloader:
            images = images.reshape(-1, config.sequence_length, config.input_size)
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, preds = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (preds == labels).sum().item()
        print(f'Test accuracy on {total} examples:{round(correct / total * 100, 2)}%')


def view():
    """可视化显示样本"""
    config = Config()
    model_path = config.model_path
    device = config.device_cpu
    model = BiRNN(config.input_size, config.hidden_size, config.num_layers, config.num_classes, device=device).to(device)
    if os.path.isfile(model_path):
        print('load model from: ', model_path)
        model.load_state_dict(torch.load(model_path))

    model.eval()
    testloader = Mnist.testloader
    images, labels = next(iter(testloader))
    with torch.no_grad():
        images = images.reshape(-1, config.sequence_length, config.input_size)
        outputs = model(images)
        _, preds = torch.max(outputs, 1)
    plt.figure()
    for i, image in enumerate(images.numpy()[:16, :, :]):
        plt.subplot(4, 4, i + 1)
        plt.imshow(image.squeeze())
        plt.title(f'label:{labels[i].item()},pred:{preds[i].item()}')
        plt.axis('off')
    plt.show()


if __name__ == '__main__':
    train()
    test()
    view()
